﻿using System;
using System.IO;
using System.Security.Cryptography;
using ooWii.Security;

namespace ooWii.IO
{
	public class WiiPartitionStream : Stream
	{
		#region WiiPartitionStream

		Stream iso;
		long dataOffset;
		byte[] decryptedData;
		int lastCluster = -1;
		long decryptedSize;
		long position;

		const int clusterSize = 0x8000;
		const int shaBytesPerCluster = 0x400;
		const int dataBytesPerCluster = 0x7c00;
		const int blockSize = 0x400;
		const int h1Index = 0x280;
		const int h2Index = 0x340;
		const int shaSize = 20;

		Ticket ticket;
		byte[] h3Table;
		SHA1Engine sha1Engine;
		AESEngine aesEngine;

		public WiiPartitionStream( Stream iso, long partitionOffset, long decryptedSize, Ticket ticket, byte[] h3Table, SHA1Engine sha1Engine, AESEngine aesEngine )
		{
			this.iso = iso;
			this.dataOffset = partitionOffset;
			this.decryptedSize = decryptedSize;
			this.ticket = ticket;
			this.h3Table = h3Table;
			this.sha1Engine = sha1Engine;
			this.aesEngine = aesEngine;
		}

		#endregion

		#region Stream

		public override bool CanRead
		{
			get
			{
				return true;
			}
		}

		public override bool CanSeek
		{
			get
			{
				return true;
			}
		}

		public override bool CanWrite
		{
			get
			{
				return false;
			}
		}

		public override void Flush()
		{
		}

		public override long Length
		{
			get
			{
				return decryptedSize;
			}
		}

		public override long Position
		{
			get
			{
				return position;
			}
			set
			{
				if( value < 0 )
				{
					throw new ArgumentOutOfRangeException( "value", "Non-negative number required." );
				}

				position = value;
			}
		}

		public override int Read( byte[] buffer, int offset, int count )
		{
			if( buffer == null )
			{
				throw new ArgumentNullException( "buffer" );
			}

			if( buffer.Length < offset + count )
			{
				throw new ArgumentException();
			}

			int read = 0;

			if( position < decryptedSize )
			{
				long left = decryptedSize - position;
				if( count > left )
				{
					count = (int)left;
				}

				while( read < count )
				{
					int clusterNumber = (int)( ( position + read ) / dataBytesPerCluster );
					int clusterIndex = (int)( ( position + read ) % dataBytesPerCluster );

					if( clusterNumber != lastCluster )
					{
						lastCluster = -1;

						byte[] encryptedBuffer = new byte[ clusterSize ];

						iso.Position = dataOffset + ( (long)clusterNumber * clusterSize );
						if( iso.Read( encryptedBuffer, 0, encryptedBuffer.Length ) != encryptedBuffer.Length )
						{
							throw new WiiException( "Cannot read cluster" );
						}

						byte[] iv = new byte[ 16 ];

						ICryptoTransform shaTransform = aesEngine.AES.CreateDecryptor( ticket.TitleKey, iv );

						try
						{
							byte[] sha1 = shaTransform.TransformFinalBlock( encryptedBuffer, 0, shaBytesPerCluster );

							byte[] h3sha1 = sha1Engine.SHA1.ComputeHash( sha1, h2Index, 8 * shaSize );
							if( !SHA1Engine.CompareSHA1( h3sha1, 0, h3Table, ( clusterNumber / 64 ) * shaSize ) )
							{
								throw new IOException( string.Format( "h3 error in cluster {0}", clusterNumber ) );
							}

							byte[] h2sha1 = sha1Engine.SHA1.ComputeHash( sha1, h1Index, 8 * shaSize );
							if( !SHA1Engine.CompareSHA1( h2sha1, 0, sha1, h2Index + ( ( ( clusterNumber / 8 ) % 8 ) * shaSize ) ) )
							{
								throw new IOException( string.Format( "h2 error in cluster {0}", clusterNumber ) );
							}

							byte[] h1sha1 = sha1Engine.SHA1.ComputeHash( sha1, 0, 31 * shaSize );
							if( !SHA1Engine.CompareSHA1( h1sha1, 0, sha1, h1Index + ( ( clusterNumber % 8 ) * shaSize ) ) )
							{
								throw new IOException( string.Format( "h1 error in cluster {0}", clusterNumber ) );
							}

							Buffer.BlockCopy( encryptedBuffer, 0x3d0, iv, 0, iv.Length );

							ICryptoTransform dataTransform = aesEngine.AES.CreateDecryptor( ticket.TitleKey, iv );

							try
							{
								decryptedData = dataTransform.TransformFinalBlock( encryptedBuffer, shaBytesPerCluster, encryptedBuffer.Length - shaBytesPerCluster );
							}
							finally
							{
								dataTransform.Dispose();
							}

							for( int i = 0; i < 31; i++ )
							{
								byte[] h0sha1 = sha1Engine.SHA1.ComputeHash( decryptedData, i * blockSize, blockSize );
								if( !SHA1Engine.CompareSHA1( h0sha1, 0, sha1, i * shaSize ) )
								{
									throw new IOException( string.Format( "h0 error in cluster {0} block {1}", clusterNumber, i ) );
								}
							}
						}
						finally
						{
							shaTransform.Dispose();
						}

						lastCluster = clusterNumber;
					}

					int bytes = dataBytesPerCluster - clusterIndex;
					if( bytes > count - read )
					{
						bytes = count - read;
					}

					Buffer.BlockCopy( decryptedData, clusterIndex, buffer, offset + read, bytes );

					read += bytes;
				}
			}

			position += read;

			return read;
		}

		public override long Seek( long offset, SeekOrigin origin )
		{
			long newPosition;

			switch( origin )
			{
			case SeekOrigin.Begin:
				newPosition = offset;
				break;

			case SeekOrigin.Current:
				newPosition = this.Position + offset;
				break;

			case SeekOrigin.End:
				newPosition = this.Length + offset;
				break;

			default:
				throw new ArgumentException( "origin" );
			}

			if( newPosition < 0 )
			{
				throw new IOException( "An attempt was made to move the file pointer before the beginning of the file." );
			}

			position = newPosition;
			return position;
		}

		public override void SetLength( long value )
		{
			throw new NotSupportedException();
		}

		public override void Write( byte[] buffer, int offset, int count )
		{
			throw new NotSupportedException();
		}

		#endregion
	}
}
